{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transfer Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the high level transfer learning APIs, you can easily customize pretrained models for feature extraction or fine-tuning. \n",
    "\n",
    "In this notebook, we will use a pre-trained Inception_V1 model. But we will operate on the pre-trained model to freeze first few layers, replace the classifier on the top, then fine tune the whole model. And we use the fine-tuned model to solve the dogs-vs-cats classification problem,"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Get the dogs-vs-cats datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the training dataset from https://www.kaggle.com/c/dogs-vs-cats and extract it. \n",
    "\n",
    "The following commands copy about 1100 images of cats and dogs into demo/cats and demo/dogs separately. \n",
    "```shell\n",
    "mkdir -p demo/dogs\n",
    "mkdir -p demo/cats\n",
    "cp train/cat.7* demo/cats\n",
    "cp train/dog.7* demo/dogs```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Get the pre-trained Inception-V1 model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the pre-trained Inception-V1 model from [Zoo](https://s3-ap-southeast-1.amazonaws.com/bigdl-models/imageclassification/imagenet/bigdl_inception-v1_imagenet_0.4.0.model) \n",
    " Alternatively, user may also download pre-trained caffe/Tensorflow/keras model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "from bigdl.dllib.nn.criterion import CrossEntropyCriterion\n",
    "from pyspark.ml import Pipeline\n",
    "from pyspark.sql.functions import col, udf\n",
    "from pyspark.sql.types import DoubleType, StringType\n",
    "\n",
    "from bigdl.dllib.nncontext import *\n",
    "from bigdl.dllib.feature.image import *\n",
    "from bigdl.dllib.keras.layers import Dense, Input, Flatten\n",
    "from bigdl.dllib.keras.models import *\n",
    "from bigdl.dllib.net import *\n",
    "from bigdl.dllib.nnframes import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc = init_nncontext(\"ImageTransferLearningExample\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "manually set model_path and image_path for training\n",
    "\n",
    "1. model_path = path to the pre-trained models. (E.g. path/to/model/bigdl_inception-v1_imagenet_0.4.0.model)\n",
    "\n",
    "2. image_path = path to the folder of the training images. (E.g. path/to/data/dogs-vs-cats/demo/\\*/\\*)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model_path = \"path/to/model/bigdl_inception-v1_imagenet_0.4.0.model\"\n",
    "image_path = \"file://path/to/data/dogs-vs-cats/demo/*/*\"\n",
    "imageDF = NNImageReader.readImages(image_path, sc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "getName = udf(lambda row:\n",
    "                  re.search(r'(cat|dog)\\.([\\d]*)\\.jpg', row[0], re.IGNORECASE).group(0),\n",
    "                  StringType())\n",
    "getLabel = udf(lambda name: 1.0 if name.startswith('cat') else 2.0, DoubleType())\n",
    "\n",
    "labelDF = imageDF.withColumn(\"name\", getName(col(\"image\"))) \\\n",
    "        .withColumn(\"label\", getLabel(col('name')))\n",
    "(trainingDF, validationDF) = labelDF.randomSplit([0.9, 0.1])\n",
    "labelDF.select(\"name\",\"label\").show(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-tune a pre-trained model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We fine-tune a pre-trained model by removing the last few layers, freezing the first few layers, and adding some new layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = ChainedPreprocessing(\n",
    "        [RowToImageFeature(), ImageResize(256, 256), ImageCenterCrop(224, 224),\n",
    "         ImageChannelNormalize(123.0, 117.0, 104.0), ImageMatToTensor(), ImageFeatureToTensor()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load a pre-trained model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the Net API to load a pre-trained model, including models saved by Analytics Zoo, BigDL, Torch, Caffe and Tensorflow. Please refer to [Net API Guide](https://analytics-zoo.github.io/master/#APIGuide/PipelineAPI/net/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_model = Net.load_bigdl(model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remove the last few layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we print all the model layers and you can choose which layer(s) to remove.\n",
    "\n",
    "When a model is loaded using Net, we can use the newGraph(output) api to define a Model with the output specified by the parameter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for layer in full_model.layers:\n",
    "    print (layer.name())\n",
    "model = full_model.new_graph([\"pool5/drop_7x7_s1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The returning model's output layer is \"pool5/drop_7x7_s1\"."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Freeze some layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We freeze layers from input to pool4/3x3_s2 inclusive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.freeze_up_to([\"pool4/3x3_s2\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add a few new layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputNode = Input(name=\"input\", shape=(3, 224, 224))\n",
    "inception = model.to_keras()(inputNode)\n",
    "flatten = Flatten()(inception)\n",
    "logits = Dense(2)(flatten)\n",
    "lrModel = Model(inputNode, logits)\n",
    "classifier = NNClassifier(lrModel, CrossEntropyCriterion(), transformer) \\\n",
    "        .setLearningRate(0.003).setBatchSize(4).setMaxEpoch(1).setFeaturesCol(\"image\") \\\n",
    "        .setCachingSample(False)\n",
    "pipeline = Pipeline(stages=[classifier])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The transfer learning can finish in a few minutes. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "catdogModel = pipeline.fit(trainingDF)\n",
    "predictionDF = catdogModel.transform(validationDF).cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictionDF.select(\"name\",\"label\",\"prediction\").sort(\"label\", ascending=False).show(10)\n",
    "predictionDF.select(\"name\",\"label\",\"prediction\").show(10)\n",
    "correct = predictionDF.filter(\"label=prediction\").count()\n",
    "overall = predictionDF.count()\n",
    "accuracy = correct * 1.0 / overall\n",
    "print(\"Test Error = %g \" % (1.0 - accuracy))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the model from transfer learning can achieve over 95% accuracy on the validation set."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We randomly select some images to show, and print the prediction results here. \n",
    "\n",
    "cat: prediction = 1.0\n",
    "dog: prediction = 2.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samplecat=predictionDF.filter(predictionDF.prediction==1.0).limit(3).collect()\n",
    "sampledog=predictionDF.filter(predictionDF.prediction==2.0).sort(\"label\", ascending=False).limit(3).collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "for cat in samplecat:\n",
    "    print (\"prediction:\"), cat.prediction\n",
    "    display(Image(cat.image.origin[5:]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for dog in sampledog:\n",
    "    print (\"prediction:\"), dog.prediction\n",
    "    display(Image(dog.image.origin[5:]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
